package com.sunpy.commonlock.zookeeper;

import com.sunpy.commonlock.pub.ZookeeperLockConfig;
import com.sunpy.commonlock.quartz.CronUtils;
import com.sunpy.commonlock.quartz.IQuartzService;
import lombok.AllArgsConstructor;
import org.apache.commons.lang3.StringUtils;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;
import org.quartz.Job;
import org.quartz.JobDataMap;
import org.quartz.JobDetail;
import org.quartz.JobExecutionContext;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.nio.charset.StandardCharsets;
import java.time.LocalDateTime;
import java.util.*;
import java.util.concurrent.CountDownLatch;

/**
 * zookeeper实现分布式锁
 * @author sunpeiyu
 * @date 2023-04-10
 */
@Component
public class ZkLockServiceImpl implements IZkLockService {

    // 控制线程执行顺序
    private final CountDownLatch cdl = new CountDownLatch(1);
    @Autowired
    private ZookeeperLockConfig config;
    @Autowired
    private IQuartzService quartzService;

    /**
     * 获取ZooKeeper实例
     * @return
     */
    @Override
    public ZooKeeper getZooKeeperInstance() {
        // ZooKeeper地址
        String url = config.getZkIp() + ":" + config.getZkPort();
        // 会话时长
        int sessionTimeout = config.getSessionTimeout();

        try {
            ZooKeeper zookeeper = new ZooKeeper(url, sessionTimeout, watchedEvent -> {});
            this.createRootNode(zookeeper);
            return zookeeper;
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 删除指定路径上的节点
     * @param zookeeper
     * @param currentPath
     * @param jobName
     */
    @Override
    public void releaseLock(ZooKeeper zookeeper, String currentPath, String jobName) {
        try {
            zookeeper.delete(currentPath, -1);
            zookeeper.close();
            // 删除定时任务
            quartzService.deleteJob(jobName);
        } catch (InterruptedException | KeeperException e) {
            throw new RuntimeException(e.getMessage());
        }
    }

    /**
     * 续期
     * 如果节点已经删除，那么续期将无用
     * @param zooKeeper
     * @param jobName
     */
    @Override
    public void renewLock(ZooKeeper zooKeeper, String jobName) {
        JobDetail jobDetail = quartzService.getDetail(jobName, TimoutJob.class);
        String currentLockPath = jobDetail.getJobDataMap().getString("currentLockPath");

        try {
            Stat stat = zooKeeper.exists(currentLockPath, false);

            if (stat != null) {
                LocalDateTime ldt = (LocalDateTime) jobDetail.getJobDataMap().get("cronTime");
                LocalDateTime updateLdt = ldt.plusSeconds(config.getRetime());
                String cron = CronUtils.convertTimePointToCron(updateLdt);
                jobDetail.getJobDataMap().put("cronTime", updateLdt);
                quartzService.updateJob(jobName, TimoutJob.class, cron, jobDetail.getJobDataMap());
            }
        } catch (KeeperException | InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 获取锁
     * @param zooKeeper
     * @param pwd
     * @param jobName
     */
    @Override
    public void getLock(ZooKeeper zooKeeper, String pwd, String jobName) {
        ZooKeeperLockInfo zooKeeperLockInfo = this.tryLock(zooKeeper, pwd);

        if (!zooKeeperLockInfo.getLockFlag()) {
            this.waitLock(zooKeeper, zooKeeperLockInfo.getPrevLockPath());
            // 当前节点删除节点释放锁之后，其他线程进来获取也得再次排队获取，拿不到继续阻塞
            this.getLock(zooKeeper, pwd, jobName);
        }
        // 超时释放锁
        this.releaseByTimeout(zooKeeper, zooKeeperLockInfo.getCurrentLockPath(), jobName);
    }

    /**
     * 尝试获取锁
     * @param zooKeeper
     * @return
     */
    private ZooKeeperLockInfo tryLock(ZooKeeper zooKeeper, String pwd) {
        if (Objects.isNull(zooKeeper)) {
            throw new RuntimeException("连接不能为空！");
        }

        String rootPath = config.getRootPath();
        String nodePathPrefix = config.getNodePathPrefix();
        String lockPath = rootPath + nodePathPrefix;
        String lockPwd = UUID.randomUUID().toString().replace("-", "");

        try {
            // 创建临时顺序节点，断开连接时，自动删除节点
            String currentLockPath = zooKeeper.create(lockPath,
                    lockPwd.getBytes(StandardCharsets.UTF_8),
                    ZooDefs.Ids.OPEN_ACL_UNSAFE,
                    CreateMode.EPHEMERAL_SEQUENTIAL);
            List<String> list = zooKeeper.getChildren(rootPath, false);
            Collections.sort(list);

            // 锁重入逻辑
            if (config.getReentry() && StringUtils.isNotEmpty(pwd)) {
                byte[] arr = zooKeeper.getData(rootPath + "/" + list.get(0), false, null);

                if (pwd.equals(new String(arr))) {
                    ZooKeeperLockInfo zooKeeperLockInfo = new ZooKeeperLockInfo();
                    zooKeeperLockInfo.setLockFlag(true);
                    zooKeeperLockInfo.setLockPwd(lockPwd);
                    zooKeeperLockInfo.setCurrentLockPath(currentLockPath);
                    return zooKeeperLockInfo;
                }
            }

            if (currentLockPath.equals(rootPath + "/" + list.get(0))) {
                ZooKeeperLockInfo zooKeeperLockInfo = new ZooKeeperLockInfo();
                zooKeeperLockInfo.setLockFlag(true);
                zooKeeperLockInfo.setLockPwd(lockPwd);
                zooKeeperLockInfo.setCurrentLockPath(currentLockPath);
                return zooKeeperLockInfo;
            } else {
                // 当前节点没有获取到锁
                String partLockPath = currentLockPath.substring(
                        currentLockPath.lastIndexOf("/") + 1);
                String prevNode = list.get(list.indexOf(partLockPath) - 1);
                String prevPath = rootPath + "/" + prevNode;
                ZooKeeperLockInfo zooKeeperLockInfo = new ZooKeeperLockInfo();
                zooKeeperLockInfo.setLockFlag(false);
                zooKeeperLockInfo.setLockPwd(lockPwd);
                zooKeeperLockInfo.setCurrentLockPath(currentLockPath);
                zooKeeperLockInfo.setPrevLockPath(prevPath);
                return zooKeeperLockInfo;
            }
        } catch (KeeperException | InterruptedException e) {
            throw new RuntimeException(e.getMessage());
        }
    }

    /**
     * 阻塞等待锁
     * @param zooKeeper
     * @param prevLockPath
     */
    private void waitLock(ZooKeeper zooKeeper, String prevLockPath) {
        try {
            Stat stat = zooKeeper.exists(prevLockPath, watchedEvent -> {
                if (Watcher.Event.EventType.NodeDeleted == watchedEvent.getType()) {
                    cdl.countDown();
                }
            });

            if (stat != null) {
                cdl.await();
            }
        } catch (KeeperException | InterruptedException e) {
            throw new RuntimeException(e.getMessage());
        }
    }

    /**
     * 超时释放锁
     * @param zooKeeper
     * @param currentLockPath
     * @param jobName
     */
    private void releaseByTimeout(ZooKeeper zooKeeper, String currentLockPath, String jobName) {
        // String jobName = "Timeout" + UUID.randomUUID().toString().replace("-","");
        Integer timeout = config.getTimeout();
        LocalDateTime ldt = LocalDateTime.now().plusSeconds(timeout);
        String cron = CronUtils.convertTimePointToCron(ldt);
        Map<String, Object> paramMap = new HashMap<>();
        paramMap.put("zooKeeper", zooKeeper);
        paramMap.put("currentLockPath", currentLockPath);
        paramMap.put("jobName", jobName);
        paramMap.put("cronTime", ldt);
        quartzService.doScheduleJob(jobName, TimoutJob.class, cron, paramMap);
    }

    /**
     * 创建根节点
     * @param zooKeeper
     * @return
     */
    private String createRootNode(ZooKeeper zooKeeper) {
        String rootPath = config.getRootPath();
        Stat stat = null;

        try {
            stat = zooKeeper.exists(rootPath, false);

            if (stat != null) {
                return rootPath;
            } else {
                return zooKeeper.create(rootPath,
                        rootPath.getBytes("UTF-8"),
                        ZooDefs.Ids.OPEN_ACL_UNSAFE,
                        CreateMode.PERSISTENT);
            }
        } catch (KeeperException | InterruptedException | UnsupportedEncodingException e) {
            throw new RuntimeException(e.getMessage());
        }
    }

    /**
     * 超时任务内部类
     */
    @AllArgsConstructor
    private final class TimoutJob implements Job {

        @Override
        public void execute(JobExecutionContext jobExecutionContext) {
            JobDataMap jobDataMap = jobExecutionContext.getJobDetail().getJobDataMap();
            ZooKeeper zooKeeper = (ZooKeeper) jobDataMap.get("zooKeeper");
            String currentLockPath = jobDataMap.getString("currentLockPath");
            String jobName = jobDataMap.getString("jobName");
            releaseLock(zooKeeper, currentLockPath, jobName);
        }
    }
}
